package com.octopus.core.limiter;

import com.octopus.core.Request;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.HashMap;
import java.util.Map;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

/**
 * 目标网站主机名绑定速率计划器
 *
 * <p>不同的目标网站域名采取不同的速率限制
 *
 * @author yang_shoulai@hoperun.com
 * @date 2019/02/28
 */
@Slf4j
@Getter
public class HostBindRateLimitPlanner implements RateLimitPlanner {

  /** 限制器集合 */
  private Map<String, RateLimiter> limiters = new HashMap<>();

  public HostBindRateLimitPlanner() {}

  /**
   * 添加绑定速率限制
   *
   * @param host 目标网站域名
   * @param rateLimiter 速率限制器
   * @return 当前速率计划器
   */
  public HostBindRateLimitPlanner addLimiter(
      @NonNull String host, @NonNull RateLimiter rateLimiter) {
    this.limiters.put(host.toLowerCase(), rateLimiter);
    return this;
  }

  @Override
  public RateLimiter getRateLimiter(Request request) {
    try {
      URL url = new URL(request.url());
      return this.limiters.get(url.getHost().toLowerCase());
    } catch (MalformedURLException e) {
      e.printStackTrace();
    }
    return null;
  }

  /** 爬虫引擎启动时启动许可信号量释放 */
  @Override
  public void onStart() {
    this.limiters.entrySet().forEach(entry -> entry.getValue().start());
  }

  /** 爬虫引擎启动时关闭时停止许可信号量释放 */
  @Override
  public void onStop() {
    this.limiters.entrySet().forEach(entry -> entry.getValue().stop());
  }
}
